#!/usr/bin/env python3
"""
视觉+Coze AI 集成节点
订阅摄像头图像，接收语音命令，调用Coze API进行图像识别
"""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
from std_msgs.msg import String
import cv2
import numpy as np
import asyncio
import threading
from typing import Optional
import secrets
from datetime import datetime
import os

from cozepy import (
    COZE_CN_BASE_URL,
    Coze,
    TokenAuth,
    Message,
    MessageObjectString,
    ChatEventType,
)


class VisionCozeNode(Node):
    """视觉+Coze AI集成节点"""
    
    def __init__(self):
        super().__init__('vision_coze_node')
        
        # 声明参数
        self.declare_parameter('coze_api_key', '')
        self.declare_parameter('coze_bot_id', '')
        self.declare_parameter('enable_auto_save', False)
        self.declare_parameter('save_path', '/tmp/echokitbot_images')
        
        # 获取参数：优先从环境变量读取，其次从ROS参数
        self.coze_api_key = os.getenv('COZE_API_KEY') or os.getenv('coze_api_key') or self.get_parameter('coze_api_key').value
        self.coze_bot_id = os.getenv('COZE_BOT_ID') or os.getenv('coze_bot_id') or self.get_parameter('coze_bot_id').value
        
        if not self.coze_api_key or not self.coze_bot_id:
            self.get_logger().error('请设置 coze_api_key 和 coze_bot_id')
            self.get_logger().info('方法1: 设置环境变量 export COZE_API_KEY=YOUR_KEY export COZE_BOT_ID=YOUR_BOT_ID')
            self.get_logger().info('方法2: 使用ROS参数 ros2 run echokitbot_microros_voice vision_coze_node --ros-args -p coze_api_key:=YOUR_KEY -p coze_bot_id:=YOUR_BOT_ID')
            raise ValueError('缺少必要参数')
        
        self.get_logger().info(f'✓ 已从{"环境变量" if os.getenv("COZE_API_KEY") or os.getenv("coze_api_key") else "ROS参数"}获取 API 配置')
        
        # 初始化 Coze 客户端
        try:
            self.coze_client = Coze(
                auth=TokenAuth(self.coze_api_key),
                base_url=COZE_CN_BASE_URL
            )
            self.get_logger().info('✓ Coze客户端初始化成功')
        except Exception as e:
            self.get_logger().error(f'Coze客户端初始化失败: {e}')
            raise
        
        # 当前图像
        self.current_image: Optional[np.ndarray] = None
        self.last_image_time = None
        self.image_source = None  # 记录图像来源
        self.image_lock = threading.Lock()
        
        # 订阅两个摄像头话题（压缩格式）
        # 使用通用话题名，避免嵌入特定设备名称
        self.external_camera_sub = self.create_subscription(
            CompressedImage,
            '/camera_raw',
            lambda msg: self.image_callback(msg, 'external_camera'),
            10
        )
        
        self.esp32_sub = self.create_subscription(
            CompressedImage,
            'camera/image_raw/compressed',
            lambda msg: self.image_callback(msg, 'esp32'),
            10
        )
        
        self.get_logger().info('等待来自以下话题的图像:')
        self.get_logger().info('  - /camera_raw (外部摄像头)')
        self.get_logger().info('  - camera/image_raw/compressed (ESP32摄像头)')
        
        # 订阅语音命令
        self.voice_command_sub = self.create_subscription(
            String,
            '/voice_command',
            self.voice_command_callback,
            10
        )
        
        # 发布识别结果
        self.result_pub = self.create_publisher(String, '/vision_result', 10)
        
        # 统计信息
        self.image_count = 0
        self.query_count = 0
        self.query_times = []
        self.max_query_history = 10
        
        # 并发控制
        self.processing_query = False
        self.query_lock = threading.Lock()
        
        self.get_logger().info(f'视觉+Coze节点已启动')
        self.get_logger().info(f'订阅摄像头话题: /camera_raw, camera/image_raw/compressed')
        self.get_logger().info(f'等待图像和语音命令...')
    
    def image_callback(self, msg: CompressedImage, source: str):
        """接收摄像头压缩图像"""
        try:
            # 解码JPEG图像
            np_arr = np.frombuffer(msg.data, np.uint8)
            cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            
            # 检查图像有效性
            if cv_image is None or cv_image.size == 0:
                self.get_logger().warn(f'收到无效图像 (来源: {source})')
                return
            
            # 检查图像尺寸
            if cv_image.shape[0] < 100 or cv_image.shape[1] < 100:
                self.get_logger().warn(f'图像尺寸过小: {cv_image.shape} (来源: {source})')
                return
            
            with self.image_lock:
                self.current_image = cv_image
                self.last_image_time = datetime.now()
                self.image_source = source
                self.image_count += 1
            
            # 每100帧打印一次日志
            if self.image_count % 100 == 0:
                self.get_logger().info(f'已接收 {self.image_count} 帧图像 (当前来源: {source})')
                
        except Exception as e:
            self.get_logger().error(f'图像解码失败 (来源: {source}): {e}')
    
    def voice_command_callback(self, msg: String):
        """接收语音命令"""
        command = msg.data.strip()
        self.get_logger().info(f'收到语音命令: {command}')
        
        # 检查是否是视觉相关命令
        vision_keywords = ['看', '识别', '图片', '照片', '拍照', '这是什么', '描述']
        video_keywords = ['视频', '录像', '连续', '动作', '过程']
        
        if any(keyword in command for keyword in video_keywords):
            # 视频分析命令（多帧）
            self.get_logger().info('检测到视频分析命令，将采集多帧图像')
            thread = threading.Thread(
                target=self._process_video_query_sync,
                args=(command,),
                daemon=True
            )
            thread.start()
        elif any(keyword in command for keyword in vision_keywords):
            # 单帧图像识别命令
            self.get_logger().info('检测到图像识别命令，将使用当前帧')
            thread = threading.Thread(
                target=self._process_vision_query_sync,
                args=(command,),
                daemon=True
            )
            thread.start()
        else:
            self.get_logger().info('命令不包含视觉关键词，跳过')
    
    def _process_vision_query_sync(self, query: str):
        """同步处理视觉查询（在独立线程中运行）- 单帧图像"""
        try:
            # 创建新的事件循环
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # 运行异步查询
            result = loop.run_until_complete(self.process_vision_query(query))
            
            loop.close()
            
        except Exception as e:
            self.get_logger().error(f'视觉查询处理失败: {e}')
    
    def _process_video_query_sync(self, query: str):
        """同步处理视频查询（在独立线程中运行）- 多帧图像"""
        try:
            # 创建新的事件循环
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # 运行异步查询
            result = loop.run_until_complete(self.process_video_query(query))
            
            loop.close()
            
        except Exception as e:
            self.get_logger().error(f'视频查询处理失败: {e}')
    
    async def process_vision_query(self, query: str) -> str:
        """
        处理视觉查询 - 单帧图像识别
        
        注意：此方法只识别当前最新的一帧图像，不是视频流分析
        """
        self.query_count += 1
        self.get_logger().info(f'开始处理视觉查询 #{self.query_count}: {query}')
        self.get_logger().info('模式: 单帧图像识别')
        
        # 获取当前图像
        with self.image_lock:
            if self.current_image is None:
                error_msg = '错误：尚未接收到摄像头图像'
                self.get_logger().error(error_msg)
                self._publish_result(error_msg)
                return error_msg
            
            image = self.current_image.copy()
            image_time = self.last_image_time
            source = self.image_source
        
        self.get_logger().info(f'使用图像时间: {image_time}, 来源: {source}')
        
        try:
            # 保存图像到临时文件
            import tempfile
            import os
            
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
                temp_path = tmp_file.name
                cv2.imwrite(temp_path, image)
            
            self.get_logger().info(f'图像已保存到: {temp_path}')
            
            try:
                # 上传图像到 Coze
                self.get_logger().info('正在上传图像到Coze...')
                image_file = self.coze_client.files.upload(file=temp_path)
                self.get_logger().info(f'✓ 图像上传成功，文件ID: {image_file.id}')
                
                # 构建消息
                message_objects = [
                    MessageObjectString.build_image(file_id=image_file.id),
                    MessageObjectString.build_text(text=query)
                ]
                
                # 发起对话
                self.get_logger().info('正在调用Coze API进行图像识别...')
                stream = self.coze_client.chat.stream(
                    bot_id=self.coze_bot_id,
                    user_id=secrets.token_urlsafe(),
                    additional_messages=[
                        Message.build_user_question_objects(message_objects)
                    ],
                )
                
                # 收集响应
                response_text = ""
                token_count = 0
                for event in stream:
                    if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                        response_text += event.message.content
                    elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                        # 记录token使用情况
                        if hasattr(event.chat, 'usage') and hasattr(event.chat.usage, 'token_count'):
                            token_count = event.chat.usage.token_count
                            self.get_logger().info(f'Token使用量: {token_count}')
                
                if not response_text:
                    response_text = "抱歉，我无法识别这张图片。"
                
                self.get_logger().info(f'✓ 识别完成，结果长度: {len(response_text)} 字符')
                self.get_logger().info(f'识别结果: {response_text[:100]}...')
                
                # 发布结果
                self._publish_result(response_text)
                
                return response_text
                
            finally:
                # 清理临时文件
                if os.path.exists(temp_path):
                    os.unlink(temp_path)
                    
        except Exception as e:
            error_msg = f'视觉查询失败: {e}'
            self.get_logger().error(error_msg)
            self._publish_result(error_msg)
            return error_msg
    
    def _publish_result(self, result: str):
        """发布识别结果"""
        msg = String()
        msg.data = result
        self.result_pub.publish(msg)
        self.get_logger().info('识别结果已发布到 /vision_result')
    
    async def process_video_query(self, query: str, duration: int = 5, interval: float = 1.0) -> str:
        """
        处理视频查询 - 多帧图像分析
        
        Args:
            query: 查询文本
            duration: 采集时长（秒）
            interval: 采样间隔（秒）
        
        Returns:
            综合分析结果
        
        注意：由于 Coze API 不直接支持视频，此方法采集多帧图像分别识别后综合
        """
        self.query_count += 1
        self.get_logger().info(f'开始处理视频查询 #{self.query_count}: {query}')
        self.get_logger().info(f'模式: 多帧采样分析 (时长: {duration}秒, 间隔: {interval}秒)')
        
        # 采集多帧图像
        frames = []
        num_frames = int(duration / interval)
        
        self.get_logger().info(f'开始采集 {num_frames} 帧图像...')
        
        for i in range(num_frames):
            with self.image_lock:
                if self.current_image is not None:
                    frames.append(self.current_image.copy())
                    self.get_logger().info(f'已采集第 {i+1}/{num_frames} 帧')
                else:
                    self.get_logger().warn(f'第 {i+1} 帧图像不可用')
            
            if i < num_frames - 1:  # 最后一帧不需要等待
                await asyncio.sleep(interval)
        
        if not frames:
            error_msg = '错误：未能采集到任何图像帧'
            self.get_logger().error(error_msg)
            self._publish_result(error_msg)
            return error_msg
        
        self.get_logger().info(f'✓ 成功采集 {len(frames)} 帧图像')
        
        # 对每一帧进行识别
        results = []
        import tempfile
        import os
        
        for idx, frame in enumerate(frames):
            try:
                self.get_logger().info(f'正在识别第 {idx+1}/{len(frames)} 帧...')
                
                # 保存到临时文件
                with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
                    temp_path = tmp_file.name
                    cv2.imwrite(temp_path, frame)
                
                try:
                    # 上传并识别
                    image_file = self.coze_client.files.upload(file=temp_path)
                    
                    message_objects = [
                        MessageObjectString.build_image(file_id=image_file.id),
                        MessageObjectString.build_text(
                            text=f"这是第{idx+1}帧图像（共{len(frames)}帧）。{query}"
                        )
                    ]
                    
                    stream = self.coze_client.chat.stream(
                        bot_id=self.coze_bot_id,
                        user_id=secrets.token_urlsafe(),
                        additional_messages=[
                            Message.build_user_question_objects(message_objects)
                        ],
                    )
                    
                    response_text = ""
                    for event in stream:
                        if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                            response_text += event.message.content
                        elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                            if hasattr(event.chat, 'usage') and hasattr(event.chat.usage, 'token_count'):
                                self.get_logger().debug(f'第{idx+1}帧 Token使用: {event.chat.usage.token_count}')
                    
                    results.append(f"第{idx+1}帧: {response_text}")
                    self.get_logger().info(f'✓ 第 {idx+1} 帧识别完成')
                    
                finally:
                    if os.path.exists(temp_path):
                        os.unlink(temp_path)
                        
            except Exception as e:
                self.get_logger().error(f'第 {idx+1} 帧识别失败: {e}')
                results.append(f"第{idx+1}帧: 识别失败")
        
        # 综合结果
        combined_result = f"视频分析结果（{len(frames)}帧，{duration}秒）：\n\n" + "\n\n".join(results)
        
        # 可选：使用 Coze 生成总结
        try:
            self.get_logger().info('正在生成视频分析总结...')
            summary_prompt = f"请根据以下多帧图像的识别结果，生成一个连贯的视频内容描述：\n\n{combined_result}"
            
            stream = self.coze_client.chat.stream(
                bot_id=self.coze_bot_id,
                user_id=secrets.token_urlsafe(),
                additional_messages=[
                    Message.build_user_question_text(summary_prompt)
                ],
            )
            
            summary = ""
            for event in stream:
                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                    summary += event.message.content
                elif event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                    if hasattr(event.chat, 'usage') and hasattr(event.chat.usage, 'token_count'):
                        self.get_logger().info(f'总结 Token使用: {event.chat.usage.token_count}')
            
            if summary:
                combined_result = f"视频分析总结：\n{summary}\n\n详细帧分析：\n{combined_result}"
                
        except Exception as e:
            self.get_logger().warn(f'生成总结失败: {e}')
        
        self.get_logger().info('✓ 视频分析完成')
        self._publish_result(combined_result)
        
        return combined_result
    
    def trigger_manual_query(self, query: str = "请详细描述这张图片的内容"):
        """手动触发视觉查询（用于测试）- 单帧"""
        self.get_logger().info(f'手动触发视觉查询: {query}')
        thread = threading.Thread(
            target=self._process_vision_query_sync,
            args=(query,),
            daemon=True
        )
        thread.start()
    
    def trigger_manual_video_query(self, query: str = "请描述这段视频的内容", duration: int = 5):
        """手动触发视频查询（用于测试）- 多帧"""
        self.get_logger().info(f'手动触发视频查询: {query}, 时长: {duration}秒')
        thread = threading.Thread(
            target=self._process_video_query_sync,
            args=(query,),
            daemon=True
        )
        thread.start()


def main(args=None):
    rclpy.init(args=args)
    
    try:
        node = VisionCozeNode()
        
        # 创建一个定时器，用于定期检查状态
        def status_callback():
            with node.image_lock:
                source_info = f', 来源: {node.image_source}' if node.image_source else ''
            node.get_logger().info(
                f'状态 - 已接收图像: {node.image_count}, '
                f'已处理查询: {node.query_count}, '
                f'当前图像: {"有" if node.current_image is not None else "无"}{source_info}'
            )
        
        timer = node.create_timer(30.0, status_callback)  # 每30秒打印一次状态
        
        rclpy.spin(node)
        
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f'节点启动失败: {e}')
    finally:
        if 'node' in locals():
            node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
